from flask import Flask, request, jsonify
from werkzeug.utils import secure_filename
import cv2
import torch
import os
import sys
import time
from flask_cors import CORS
import numpy as np
from ultralytics.nn.tasks import DetectionModel
torch.serialization.add_safe_globals([DetectionModel])

# 添加本地YOLOv5 v6.0代码路径
yolov5_path = r'D:\competition\yolo\mayolo\yolov5-6.0-my-M-A'

sys.path.insert(0, yolov5_path)

# 导入v6.0版本专用模块
from models.experimental import attempt_load
from utils.datasets import letterbox
from utils.general import non_max_suppression, check_img_size

app = Flask(__name__)
CORS(app)  # 启用跨域请求

# 配置参数
UPLOAD_FOLDER = 'temp_uploads'
ALLOWED_EXTENSIONS = {'mp4', 'avi', 'mov'}
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = 500 * 1024 * 1024  # 500MB

# 创建临时目录
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# 硬件配置
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
HALF = DEVICE.type != 'cpu'  # 是否使用半精度

# 加载自定义模型
MODEL_PATH = 'models/best.pt'  # 修改为你的模型路径
model = attempt_load(MODEL_PATH, map_location=DEVICE)  # v6.0加载方式
stride = int(model.stride.max())  # 模型步长
imgsz = check_img_size(640, s=stride)  # 验证图片尺寸
if HALF:
    model.half()  # 转换为半精度

# 自定义类别（必须与训练时一致）

class_names = ['Apple___Apple_scab', 'Apple___Black_rot', 'Apple___Cedar_apple_rust', 'Apple___healthy', 'Blueberry___healthy', 'Cherry___Powdery_mildew', 'Cherry___healthy', 'Corn___Cercospora_leaf_spot Gray_leaf_spot', 'Corn___Common_rust', 'Corn___Northern_Leaf_Blight', 'Corn___healthy', 'Grape___Black_rot', 'Grape___Esca_(Black_Measles)', 'Grape___Leaf_blight_(Isariopsis_Leaf_Spot)', 'Grape___healthy', 'Orange___Haunglongbing_(Citrus_greening)', 'Peach___Bacterial_spot', 'Peach___healthy', 'Pepper,_bell___Bacterial_spot', 'Pepper,_bell___healthy', 'Potato___Early_blight', 'Potato___Late_blight', 'Potato___healthy', 'Raspberry___healthy', 'Soybean___healthy', 'Squash___Powdery_mildew', 'Strawberry___Leaf_scorch', 'Strawberry___healthy', 'Tomato___Bacterial_spot', 'Tomato___Early_blight', 'Tomato___Late_blight', 'Tomato___Leaf_Mold', 'Tomato___Septoria_leaf_spot', 'Tomato___Spider_mites Two-spotted_spider_mite', 'Tomato___Target_Spot', 'Tomato___Tomato_Yellow_Leaf_Curl_Virus', 'Tomato___Tomato_mosaic_virus', 'Tomato___healthy'] # 示例，请替换实际类别

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/detect', methods=['POST'])
def video_detection():
    # 检查文件上传
    if 'video' not in request.files:
        return jsonify({'error': 'No video found'}), 400
    
    file = request.files['video']
    if file.filename == '':
        return jsonify({'error': 'Empty filename'}), 400
    
    if not allowed_file(file.filename):
        return jsonify({'error': 'Invalid file type'}), 400

    # 保存临时文件
    filename = secure_filename(file.filename)
    save_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    file.save(save_path)
    
    try:
        start_time = time.time()
        
        # 获取处理参数
        conf_thres = float(request.form.get('conf_thres', 0.25))
        iou_thres = float(request.form.get('iou_thres', 0.45))
        max_frames = int(request.form.get('max_frames', 0))  # 0表示处理全部
        
        # 执行检测
        results = process_video(
            save_path,
            conf_thres=conf_thres,
            iou_thres=iou_thres,
            max_frames=max_frames
        )
        
        proc_time = round(time.time() - start_time, 2)
        
        return jsonify({
            'status': 'success',
            'processing_time': proc_time,
            'frame_count': len(results),
            'detections': results   
        })
    
    except Exception as e:
        return jsonify({'error': str(e)}), 500
    
    finally:
        # 清理临时文件
        if os.path.exists(save_path):
            os.remove(save_path)

def process_video(video_path, conf_thres=0.25, iou_thres=0.45, max_frames=0):
    """
    视频处理核心函数
    返回格式：[{
        "frame": 帧序号,
        "detections": [{
            "class": 类别ID,
            "class_name": 类别名称,
            "confidence": 置信度,
            "bbox": [x1,y1,x2,y2]
        }]
    }]
    """
    cap = cv2.VideoCapture(video_path)
    results = []
    frame_idx = 0
    
    # 获取视频参数
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    if max_frames <= 0: 
        max_frames = total_frames
    
    while cap.isOpened() and frame_idx < max_frames:
        success, frame = cap.read()
        if not success:
            break
        
        # 预处理
        img = letterbox(frame, imgsz, stride=stride, auto=False)[0]  # v6.0参数
        img = np.ascontiguousarray(img[:, :, ::-1].transpose(2, 0, 1))  # BGR转RGB，HWC转CHW
        img = torch.from_numpy(img).to(DEVICE)
        img = img.half() if HALF else img.float()  # uint8转浮点
        img /= 255.0  # 0-255转0.0-1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
        
        # 推理
        with torch.no_grad():
            try:
                raw_pred = model(img, augment=False)
                pred = raw_pred[0]  # 获取第一个输出
            except Exception as e:
                print(f"推理错误: {str(e)}")
                pred = None
        
        # NMS处理
        try:
            pred = non_max_suppression(
                pred, 
                conf_thres=conf_thres,
                iou_thres=iou_thres,
                classes=None
            ) if pred is not None else []
        except Exception as e:
            print(f"NMS错误: {str(e)}")
            pred = []
        
        # 解析结果
        frame_detections = []
        for det in pred:  # 遍历所有检测结果
            if det is not None and len(det):
                # 转换为CPU上的numpy数组
                det_np = det.cpu().numpy()
                for *xyxy, conf, cls in det_np:
                    # 验证置信度阈值
                    if conf < conf_thres:
                        continue
                    
                    # 转换坐标并确保数值有效性
                    try:
                        bbox = [round(float(x), 2) for x in xyxy]
                        class_id = int(cls)
                        frame_detections.append({
                            "class": class_id,
                            "class_name": class_names[class_id],
                            "confidence": round(float(conf), 4),
                            "bbox": bbox
                        })
                    except IndexError:
                        print(f"无效类别ID: {class_id} (总类别数: {len(class_names)})")
                    except Exception as e:
                        print(f"解析错误: {str(e)}")

        results.append({
            "frame": frame_idx,
            "time": round(frame_idx/fps, 2),
            "detections": frame_detections
        })
        
        frame_idx += 1
    
    cap.release()
    return results

def warmup():
    """GPU预热"""
    img = torch.zeros((1, 3, imgsz, imgsz), device=DEVICE)
    _ = model(img.half() if HALF else img) if DEVICE.type != 'cpu' else None

# 服务启动时预热模型
warmup()

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True, threaded=True)